import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from sklearn import svm
from sklearn.model_selection import GridSearchCV

# 载入数据
df = pd.read_csv('watermelon_3_3.csv')
x = df[['density']].values
y = df['sugar_ratio'].values

# 初始化回归器（自动选择最优核函数和参数C）
svr = GridSearchCV(
    svm.SVR(), cv=5, n_jobs=-1,
    param_grid={
        "kernel" : ['linear', 'rbf', 'poly'],
        "C" : [1e-4, 1e-3, 1e-2, 1e-1, 1, 1e1, 1e2, 1e3, 1e4]
    }
)
print('\n正在选择最优超参数，这可能需要一段时间……')
svr.fit(x, y) # 训练回归器

sv_kernel = svr.best_estimator_.kernel
sv_num = svr.best_estimator_.support_.shape[0]
print('\nSVR核函数：', sv_kernel)
print('支持向量占比：', sv_num / len(x))

# 输出预测结果
result = svr.predict(x)
print('\n密度\t原含糖率\t预测含糖率')
for i in range(x.shape[0]):
    print(x[i][0], '\t', y[i], '\t\t', result[i])

# 预测结果可视化
sv_idx = svr.best_estimator_.support_ # 支持向量索引
## 绘制样本点和支持向量
plt.scatter(x, y, label='data', zorder=1)
plt.scatter(x[sv_idx], y[sv_idx], c='r', s=50, label='SVR support vectors', zorder=2)
## 绘制预测结果
x_result = np.c_[x, result]
x_result_sorted = x_result[x_result[:, 0].argsort()]
plt.plot(x_result_sorted[:, 0], x_result_sorted[:, 1], label='SVR predict result')
## 设置标题
plt.xlabel('density')
plt.ylabel('sugar_ratio')
plt.title('SVR')
plt.legend()
plt.savefig('SVR.png')
plt.show()
